import bisect
import numpy as np


def expand_targets(targets, T=500, burnin=0):
    y = np.tile(targets.copy(), [T, 1, 1])
    y[:burnin] = 0
    return y


def one_hot(mbt, num_classes):
    out = np.zeros([mbt.shape[0], num_classes])
    out[np.arange(mbt.shape[0], dtype="int"), mbt.astype("int")] = 1
    return out


def find_first(a, tgt):
    return bisect.bisect_left(a, tgt)


def cast_evs(evs):
    ts = (evs[:, 0] * 1e6).astype("uint64")
    ad = (evs[:, 1:]).astype("uint64")
    return ts, ad


def get_subsampled_coordinates(evs, ds_h, ds_w):
    x_coords = evs[:, 1] // ds_w
    y_coords = evs[:, 2] // ds_h
    if x_coords.dtype != np.int:
        x_coords = x_coords.astype(int)
    if y_coords.dtype != np.int:
        y_coords = y_coords.astype(int)
    return x_coords, y_coords


def get_binary_frame_np(arr, evs, ds_w=1, ds_h=1):
    x_coords, y_coords = get_subsampled_coordinates(evs, ds_h, ds_w)
    arr[x_coords, y_coords] = 2 * evs[:, 3] - 1


def get_binary_frame(arr, evs, ds_w=1, ds_h=1):
    """
    get_binary_frame_np 函数以 arr、evs、ds_w 和 ds_h 为参数。
    它调用带有 evs、ds_h 和 ds_w 的 get_subsampled_coordinates 函数来获取 x_coords 和 y_coords。
    然后，它将 2 * evs[:, 3] - 1 的值赋值给 x_coords 和 y_coords 指定坐标处的 arr 元素。
    """
    x_coords, y_coords = get_subsampled_coordinates(evs, ds_h, ds_w)
    arr[x_coords, y_coords] = 1


def get_slice(times, addrs, start_time, end_time):
    try:
        idx_beg = find_first(times, start_time)
        idx_end = find_first(times[idx_beg:], end_time) + idx_beg
        return times[idx_beg:idx_end] - times[idx_beg], addrs[idx_beg:idx_end]
    except IndexError:
        raise IndexError("Empty batch found")


def get_event_slice(times, addrs, start_time, T, size=[128, 128], ds=1, dt=1000):
    """
    get_event_slice 函数将 times、addrs、start_time、T、size=[128, 128]、ds=1 和 dt=1000 作为参数。
    它会尝试查找 times 中 start_time 第一次出现的索引（使用名为 find_first 的函数）。
    然后在从 idx_beg 开始的时间子列表中查找 start_time + T * dt 第一次出现的索引。它返回的是使用从 idx_beg 到 idx_end 的时间和地址子列表以适当参数调用 chunk_evs_pol_dvs 函数的结果。
    如果在此过程中出现 IndexError，则会引发一条自定义异常消息 "Empty batch found"（发现空批次）。
    """
    try:
        idx_beg = find_first(times, start_time)
        idx_end = find_first(times[idx_beg:], start_time + T * dt) + idx_beg
        return chunk_evs_pol_dvs(
            times[idx_beg:idx_end],
            addrs[idx_beg:idx_end],
            deltat=dt,
            chunk_size=T,
            size=size,
            ds_w=ds,
            ds_h=ds,
        )
    except IndexError:
        raise IndexError("Empty batch found")


def get_tmad_slice(times, addrs, start_time, T):
    """
    get_tmad_slice 函数将 times、addrs、start_time 和 T 作为参数。
    它会尝试查找 times 中 start_time 第一次出现的索引（使用名为 find_first 的函数）。
    然后在从 idx_beg 开始的时间子列表中查找 start_time + T 第一次出现的索引。
    从 idx_beg 到 idx_end，返回时间子列表和相应地址子列表的列栈。
    如果在此过程中出现 IndexError，则会引发一条自定义异常消息 "发现空批次"。
    """
    try:
        idx_beg = find_first(times, start_time)
        idx_end = find_first(times[idx_beg:], start_time + T) + idx_beg
        return np.column_stack([times[idx_beg:idx_end], addrs[idx_beg:idx_end]])
    except IndexError:
        raise IndexError("Empty batch found")


def get_time_surface(evs, invtau=1e-6, size=(346, 260, 2)):
    tr = np.zeros(size, "int64") - np.inf

    for ev in evs:
        tr[ev[2], ev[1], ev[3]] = ev[0]

    a = np.exp(tr[:, :, 0] * invtau) - np.exp(tr[:, :, 1] * invtau)

    return a


def chunk_evs_dvs(evs, deltat=1000, chunk_size=500, size=[304, 240], ds_w=1, ds_h=1):
    t_start = evs[0, 0]
    ts = range(t_start + chunk_size, t_start + chunk_size * deltat, deltat)
    chunks = np.zeros([len(ts)] + size, dtype="int8")
    idx_start = 0
    idx_end = 0
    for i, t in enumerate(ts):
        idx_end += find_first(evs[idx_end:, 0], t)
        if idx_end > idx_start:
            get_binary_frame_np(
                chunks[i, ...], evs[idx_start:idx_end], ds_h=ds_h, ds_w=ds_w
            )
        idx_start = idx_end
    return chunks


def frame_evs(times, addrs, deltat=1000, duration=500, size=[240], downsample=[1]):
    """
    frame_evs 函数以 times、addrs、deltat、duration、size 和 downsample 为参数。
    它根据 times 中的第一个值计算 t_start，并创建一个 ts 值范围。然后返回 ts。
    """
    t_start = times[0]
    ts = range(t_start, t_start + duration * deltat, deltat)
    chunks = np.zeros([len(ts)] + size, dtype="int8")
    idx_start = 0
    idx_end = 0
    for i, t in enumerate(ts):
        idx_end += find_first(times[idx_end:], t)
        if idx_end > idx_start:
            ee = addrs[idx_start:idx_end]
            ev = [(ee[:, i] // d).astype(np.int) for i, d in enumerate(downsample)]
            np.add.at(chunks, tuple([i] + ev), 1)
        idx_start = idx_end
    return chunks


def chunk_evs_pol_dvs(
        times, addrs, deltat=1000, chunk_size=500, size=[2, 304, 240], ds_w=1, ds_h=1
):
    """
    chunk_evs_pol_dvs 函数将 times、addrs、deltat、chunk_size、size、ds_w 和 ds_h 作为参数。
    它根据 times 中的第一个值计算 t_start，并创建一个 ts 值范围。
    如果 idx_end 大于 idx_start，它就会提取 addrs 的一个子集，并对其执行一些计算。最后，它会更新 idx_start，并返回分块。
    """
    t_start = times[0]
    ts = range(t_start, t_start + chunk_size * deltat, deltat)
    chunks = np.zeros([len(ts)] + size, dtype="int8")
    idx_start = 0
    idx_end = 0
    for i, t in enumerate(ts):
        idx_end += find_first(times[idx_end:], t)
        if idx_end > idx_start:
            ee = addrs[idx_start:idx_end]
            pol, x, y = (
                ee[:, 2],
                (ee[:, 0] // ds_w).astype(np.int),
                (ee[:, 1] // ds_h).astype(np.int),
            )
            np.add.at(chunks, (i, pol, x, y), 1)
        idx_start = idx_end
    return chunks


def chunk_evs_pol_dvs_gesture(data, dt=1000, T=500, size=[2, 346, 260], ds=[4, 4]):
    """
    chunk_evs_pol_dvs_gesture 函数接收多个参数：
        数据： 代表事件数据的列表，其中每个内列表都包含值 [t、p、x、y]。
        dt：事件块之间的持续时间。
        T：要创建的事件块数量。@ data/thu_chl/process.py@func:get_eventFrame : time_num：时间数： 一个整数，指定要将事件分割成的时间窗口数量。
        size： 大小： 表示每个数据块大小的列表。
        ds： 表示 x 和 y 坐标下采样因子的列表。

    函数首先会尝试从数据列表的第一个元素中提取起始时间 t_start。如果出现异常，函数会将 t_start 设置为 0。
    然后，函数根据 t_start、T 和 dt 创建一个时间值范围 ts。
    函数初始化一个形状为 [len(ts)] + size、数据类型为 int64 的 NumPy 数组 chunks，用于存储分块的事件数据。
    然后迭代 ts 中的时间值，并执行以下步骤：
        通过查找数据中第一个大于当前时间值加上 dt 的时间值的索引，更新 idx_end 变量。
        如果 idx_end 大于 idx_start，则使用 idx_start 和 idx_end 索引从数据中提取相关事件数据 ee。
        它将 ee 的 x 和 y 坐标除以 ds 中的相应值，并对结果进行下采样。极性值按原样提取。
        根据当前的时间索引、极性和下采样的 x 和 y 坐标，通过递增相应的元素来更新块数组。
        最后，函数返回包含分块事件数据的分块数组。

    """
    try:
        t_start = data[0][0]
    except Exception as e:
        t_start = 0
    ts = range(t_start, t_start + T * dt, dt)  # T:训练分的事件段数
    chunks = np.zeros([len(ts)] + size, dtype="int64")  # (T,2,173,130) 此处已经被ds整除得到
    idx_start = 0
    idx_end = 0
    for i, t in enumerate(ts):
        idx_end += find_first(data[idx_end:, 0], t + dt)
        if idx_end > idx_start:
            ee = data[idx_start:idx_end, 1:]  # ee:[[p,x,y],[p,x,y]...]
            pol, x, y = (
                ee[:, 0],  # pol
                np.floor(ee[:, 1] / ds[0]).astype(np.int64),  # x
                np.floor(ee[:, 2] / ds[1]).astype(np.int64),  # y
            )
            np.add.at(chunks, (i, pol, x, y), 1)  # i:0-T,pol:0/1,x:0-173,y:0-130
        idx_start = idx_end
    return chunks  # chunks (T,2,346/ds[0],260/ds[1]) s.t.(T,2,173,130)


def chunk_evs_pol_dvs_gait(data, dt=1000, T=500, size=[2, 304, 240], ds=[4, 4]):
    """
    chunk_evs_pol_dvs_gait 函数接收多个参数：
        数据： 代表事件数据的列表，其中每个内列表都包含值 [t、p、x、y]。
        dt：事件块之间的持续时间。
        T：要创建的事件块数量。
        size： 大小： 表示每个数据块大小的列表。
        ds： 表示 x 和 y 坐标下采样因子的列表。

    函数首先检查数据的长度是否为 0，如果为 0，则根据 T、size 和 ds 返回一个形状为零的数组。
    如果数据不为空，函数将从数据列表的第一个元素中提取开始时间 t_start。
    然后，函数根据 t_start、T 和 dt 创建一个时间值范围 ts。
    函数初始化一个形状为 [len(ts)] + size、数据类型为 int64 的 NumPy 数组 chunks，用于存储分块的事件数据。
    然后将 idx_start 和 idx_end 变量初始化为 0。

    函数遍历 ts 中的时间值，并执行以下步骤：
        它通过查找数据中第一个大于当前时间值加上 dt 的时间值的索引来更新 idx_end 变量。
        如果 idx_end 大于 idx_start，则使用 idx_start 和 idx_end 索引从数据中提取相关事件数据 ee。
        它将 ee 的 x 和 y 坐标除以 ds 中的相应值，并对结果进行下采样。极性值按原样提取。
        根据当前的时间索引、极性和下采样的 x 和 y 坐标，通过递增相应的元素来更新块数组。
        最后，函数返回包含分块事件数据的分块数组。
    """
    if len(data) == 0:
        return np.zeros(
            [T] + [2] + [int(128 / ds[0])] + [int(128 / ds[1])], dtype="int64"
        )
    t_start = data[0][0]
    ts = range(t_start, t_start + T * dt, dt)
    chunks = np.zeros(
        [len(ts)] + [2] + [int(128 / ds[0])] + [int(128 / ds[1])], dtype="int64"
    )
    idx_start = 0
    idx_end = 0
    for i, t in enumerate(ts):
        idx_end += find_first(data[idx_end:, 0], t + dt)
        if idx_end > idx_start:
            ee = data[idx_start:idx_end, 1:]
            pol, x, y = (
                ee[:, 2],
                np.floor(ee[:, 0] / ds[0]).astype(np.int64),
                np.floor(ee[:, 1] / ds[1]).astype(np.int64),
            )
            np.add.at(chunks, (i, pol, x, y), 1)
        idx_start = idx_end
    return chunks
